import cv2
import numpy as np
from rembg import remove
from PIL import Image


class BackgroundReplacer:
    def __init__(self):
        # 预定义标准背景色（BGR格式）
        self.bg_colors = {
            'white': (255, 255, 255),
            'blue': (219, 142, 67),  # RGB(67,142,219)的BGR表示
            'red': (60, 20, 220),  # RGB(220,20,60)的BGR表示
            'gray': (128, 128, 128)
        }

    def _remove_background(self, image):
        """使用U2-Net去除背景（保持原始尺寸）"""
        pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        result = remove(pil_img)
        return cv2.cvtColor(np.array(result), cv2.COLOR_RGBA2BGRA)

    def _change_background(self, image, bg_color):
        """精确背景替换（原图尺寸）"""
        # 分离alpha通道
        alpha = image[:, :, 3] / 255.0
        alpha = cv2.merge([alpha, alpha, alpha])

        # 生成新背景
        background = np.full_like(image[:, :, :3], bg_color, dtype=np.uint8)

        # 混合计算
        foreground = image[:, :, :3].astype(float)
        background = background.astype(float)
        blended = cv2.multiply(alpha, foreground) + cv2.multiply(1.0 - alpha, background)

        return blended.astype(np.uint8)

    def process(self, input_path, output_path, bg_color_name='white'):
        """
        处理流程：
        1. 读取原始图片
        2. 去除背景（保持原尺寸）
        3. 替换背景颜色
        4. 保存结果
        """
        # 验证背景颜色
        if bg_color_name not in self.bg_colors:
            available = list(self.bg_colors.keys())
            raise ValueError(f"无效背景颜色，可选：{available}")

        # 读取图像并验证
        img = cv2.imread(input_path)
        if img is None:
            raise FileNotFoundError(f"无法读取输入文件：{input_path}")

        # 记录原始尺寸
        original_h, original_w = img.shape[:2]

        try:
            # 去除背景
            no_bg_img = self._remove_background(img)

            # 替换背景颜色
            new_bg_img = self._change_background(no_bg_img, self.bg_colors[bg_color_name])

            # 保持原始尺寸
            final_img = cv2.resize(new_bg_img, (original_w, original_h))

            # 保存结果
            cv2.imwrite(output_path, final_img)
            print(f"成功处理：{input_path} -> {output_path} | 背景色：{bg_color_name}")

        except Exception as e:
            raise RuntimeError(f"处理失败：{str(e)}")


if __name__ == "__main__":
    # 使用示例
    processor = BackgroundReplacer()

    # 替换为蓝色背景（保持原尺寸）
    processor.process(
        input_path="../../input.jpg",
        output_path="output_blue.jpg",
        bg_color_name="blue"
    )

    # 替换为白色背景（保持原尺寸）
    processor.process(
        input_path="../../input.jpg",
        output_path="output_white.jpg",
        bg_color_name="white"
    )